import os
import time
import requests
from bs4 import BeautifulSoup

# 数据初始化
class viode_ini:
    # 初始化游标
    max_cursor = 0
    # 初始化视频数量
    quantity = 0
    # 成功下载的视频数量
    success = 0
    # 失败下载的视频数量
    error = 0

    # 全局请求头
    header = {
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/101.0.0.0 Safari/537.36'
    }

    # 初始化文件夹
    def __init__(self):
        # 判断文件夹是否存在
        if not os.path.exists('video'):
            os.mkdir('video')
        if not os.path.exists('video/主页'):
            os.mkdir('video/主页')
        if not os.path.exists('video/喜欢'):
            os.mkdir('video/喜欢')
        if not os.path.exists('video/合集'):
            os.mkdir('video/合集')

    # 链接重定向
    def redirect(self, url):
        response = requests.get(url, headers=self.header)
        return response.url

    # 视频下载
    def download(self, url, title, path):
        # 视频下载
        try:
            response = requests.get(url, headers=self.header)
            with open(path + title + '.mp4', 'wb') as f:
                f.write(response.content)
            self.success += 1
        except Exception as e:
            self.error += 1

    # 结束控制台输出
    def end(self):
        print('=====================')
        print('视频总数：' + str(self.quantity))
        print('成功下载：' + str(self.success))
        print('下载失败：' + str(self.error))
        print('=====================')
        restart()

# 报错日志
def error_log(error):
    # 设置报错日志文件名
    log_name = time.strftime('%Y-%m-%d', time.localtime(time.time()))
    # 判断文件夹是否存在
    if not os.path.exists('video/log'):
        os.mkdir('video/log')
    # 获取当前时间
    now_time = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time()))
    # 日志格式
    log = f'=====================\n' \
          f'[{now_time}]\n' \
          f'{error}\n' \
          f'=====================\n'
    # 写入日志
    with open('video/log/' + log_name + '.txt', 'a', encoding='utf-8') as f:
        f.write(log)

# 特殊字符处理
def replace(title):
    title = title.replace('\\', '')
    title = title.replace('/', '')
    title = title.replace(':', '')
    title = title.replace('*', '')
    title = title.replace('?', '')
    title = title.replace('"', '')
    title = title.replace('<', '')
    title = title.replace('>', '')
    title = title.replace('|', '')
    title = title.replace('\n', '')
    return title

# 视频地址获取
def video(data):
    try:
        # 获取视频下载地址
        url = data['video']['download_addr']['url_list'][0]
        return url
    except KeyError:
        return None

# 结束判断
def restart():
    operation = input('是否继续爬取？(y/n)：')
    if operation == 'y':
        start()
    elif operation == 'n':
        exit()

# 程序入口
def start():
    print('=====================')
    print('欢迎使用抖音视频下载器')
    print('作者：y7wanghui@gmail.com只供娱乐')
    print('=====================')
    print('1.下载主页视频')
    print('2.下载喜欢视频')
    print('3.下载合集视频')
    print('4.退出')
    Type = input('请输入你的选择：')
    if Type == '1':
        # 下载主页视频
        home()
    elif Type == '2':
        # 下载喜欢视频
        like()
    elif Type == '3':
        # 下载合集视频
        collection()
    elif Type == '4':
        # 退出
        exit()
    else:
        print('输入错误，请重新输入')
        start()

# 主页视频下载
def home():
    # 初始化数据
    viode_data = viode_ini()
    # 获取用户主页链接
    url = input('请输入用户主页链接：')
    # 获取真实链接
    url = viode_data.redirect(url)
    # 获取sec_uid
    sec_uid = url.split('user/')[1].split('?')[0]
    # 开始执行任务
    while True:
        # 获取视频列表
        url = f'https://m.douyin.com/web/api/v2/aweme/post/?reflow_source=reflow_page&sec_uid={sec_uid}&count=21&max_cursor={viode_data.max_cursor}'
        response = requests.get(url, headers=viode_data.header)
        # 获取视频列表数据
        aweme_list = response.json()['aweme_list']
        # 读取视频列表数据
        for aweme in aweme_list:
            # 更新视频总数
            viode_data.quantity += 1
            # 获取视频标题
            title = aweme['desc']
            # 替换标题中的特殊字符
            title = replace(title)
            # 更新视频标题
            title = '【'+str(viode_data.quantity) + '】 ' + title
            # 获取作者名称
            author = aweme['author']['nickname']
            # 替换作者名称中的特殊字符
            author = replace(author)
            # 获取视频下载地址
            url = video(aweme)
            # 设置下载路径
            path = f'video/主页/{author}/'
            # 判断视频是否为图文视频
            if url is not None:
                # 判断文件夹是否存在
                if not os.path.exists(path):
                    # 创建文件夹
                    os.makedirs(path)
                # 下载视频
                print(f'视频正在下载：{title}')
                # 视频下载
                viode_data.download(url, title, path)
            else:
                viode_data.error += 1
                print(f'视频下载失败：{title} （图文视频，跳过下载）')
                continue
        # 判断是否还有下一页
        if response.json()['has_more']:
            # 更新游标
            viode_data.max_cursor = response.json()['max_cursor']
        else:
            # 结束控制台输出
            viode_data.end()

# 喜欢视频下载
def like():
    # 初始化数据
    viode_data = viode_ini()
    # 获取用户主页链接
    url = input('请输入用户主页链接：')
    # 获取真实链接
    url = viode_data.redirect(url)
    # 获取sec_uid
    sec_uid = url.split('user/')[1].split('?')[0]
    # 获取收藏者信息
    url = f'https://www.iesdouyin.com/web/api/v2/user/info/?sec_uid={sec_uid}'
    response = requests.get(url, headers=viode_data.header)
    # 获取收藏者名称
    nickname = response.json()['user_info']['nickname']
    # 替换收藏者名称中的特殊字符
    nickname = replace(nickname)

    # 开始执行任务
    while True:
        url = f'https://m.douyin.com/web/api/v2/aweme/like/?reflow_source=reflow_page&sec_uid={sec_uid}&count=21&max_cursor={viode_data.max_cursor}'
        response = requests.get(url, headers=viode_data.header)
        # 获取视频列表数据
        aweme_list = response.json()['aweme_list']
        # 读取视频列表数据
        for aweme in aweme_list:
            # 更新视频总数
            viode_data.quantity += 1
            # 获取视频标题
            title = aweme['desc']
            # 替换标题中的特殊字符
            title = replace(title)
            # 更新视频标题
            title = '【' + str(viode_data.quantity) + '】 ' + title
            # 获取视频下载地址
            url = video(aweme)
            # 设置下载路径
            path = f'video/喜欢/{nickname}/'
            # 判断视频是否为图文视频
            if url is not None:
                # 判断文件夹是否存在
                if not os.path.exists(path):
                    # 创建文件夹
                    os.makedirs(path)
                # 下载视频
                print(f'视频正在下载：{title}')
                # 视频下载
                viode_data.download(url, title, path)
            else:
                viode_data.error += 1
                print(f'视频下载失败：{title} （图文视频，跳过下载）')
                continue
        # 判断是否还有下一页
        if response.json()['has_more']:
            # 更新游标
            viode_data.max_cursor = response.json()['max_cursor']
        else:
            # 结束控制台输出
            viode_data.end()

# 合集视频下载
def collection():
    # 初始化数据
    viode_data = viode_ini()
    # 获取合集链接
    url = input('请输入合集链接：')
    # 获取真实链接
    url = viode_data.redirect(url)
    mix_id = url.split('detail/')[1].split('/')[0]
    # 获取合集信息
    url = f'https://www.iesdouyin.com/share/mix/detail/{mix_id}'
    response = requests.get(url, headers=viode_data.header)
    # 获取合集信息
    soup = BeautifulSoup(response.text, 'html.parser')
    # 获取合集名称
    collection_name = soup.select('span[class="mix-info-name-text"]')[0].text
    # 替换合集名称中的特殊字符
    collection_name = replace(collection_name)

    # 开始执行任务
    while True:
        url = f'https://www.iesdouyin.com/web/api/mix/item/list/?reflow_source=reflow_page&mix_id={mix_id}&count=10&cursor={viode_data.max_cursor}'
        response = requests.get(url, headers=viode_data.header)
        # 获取视频列表数据
        aweme_list = response.json()['aweme_list']
        # 读取视频列表数据
        for aweme in aweme_list:
            # 更新视频总数
            viode_data.quantity += 1
            # 获取视频标题
            title = aweme['desc']
            # 替换标题中的特殊字符
            title = replace(title)
            # 更新视频标题
            title = '【第' + str(viode_data.quantity) + '集】 ' + title
            # 获取作者名称
            nickname = aweme['author']['nickname']
            # 替换作者名称中的特殊字符
            nickname = replace(nickname)
            # 获取视频下载地址
            url = video(aweme)
            # 设置下载路径
            path = f'video/合集/{nickname}/{collection_name}'
            # 判断视频是否为图文视频
            if url is not None:
                # 判断作者文件夹是否存在
                if not os.path.exists(f'video/合集/{nickname}'):
                    # 创建作者文件夹
                    os.makedirs(f'video/合集/{nickname}')
                # 判断合集文件夹是否存在
                if not os.path.exists(path):
                    # 创建合集文件夹
                    os.makedirs(path)
                # 下载视频
                print(f'视频正在下载：{title}')
                # 视频下载
                viode_data.download(url, title, path)
            else:
                viode_data.error += 1
                print(f'视频下载失败：{title} （图文视频，跳过下载）')
                continue
        # 判断是否还有下一页
        if response.json()['has_more']:
            # 更新游标
            viode_data.max_cursor = response.json()['cursor']
        else:
            # 结束控制台输出
            viode_data.end()

if __name__ == '__main__':
    viode_ini()
    start()
